'''
Author: your name
Date: 2021-04-01 10:07:15
LastEditTime: 2021-04-01 16:50:02
LastEditors: Please set LastEditors
Description: In User Settings Edit
FilePath: \genetic-algorithm-optimization\GA.py
'''

from __future__ import division
import numpy as np
import random
import math
import copy
import matplotlib as mpl
import matplotlib.pyplot as plt
from operator import itemgetter
from globalVal import *


class GA(object):
    def __init__(self, maxiter, sizepop, lenchrom, pc, pm, Fobj):
        """
        maxiter：最大迭代次数
        sizepop：种群数量
        lenchrom：染色体长度
        pc：交叉概率
        pm：变异概率
        Fobj：价值函数
        """
        self.maxiter = maxiter
        self.sizepop = sizepop
        self.lenchrom = lenchrom
        self.pc = pc
        self.pm = pm
        self.Fobj = Fobj
        self.pop = []
        self.bestindividual = None

    # 初始化种群：返回一个二维数组，第一维是种子，第二维是编码基因
    def Initialization(self):
        population = []
        for i in range(self.sizepop):
            chromosome = []
            stepFlag = 1
            while stepFlag != 2:
                for gene in range(self.lenchrom):
                    chromosome.append(random.randint(0, 1))
                stepFlag, scores = self.Fobj.step_online(chromosome)
            fitness = np.mean(scores)/10000
            population.append({'Chromosome':chromosome, 'fitness':fitness})
        return population

    def selectBest(self, pop):
        s_inds = sorted(pop, key = itemgetter("fitness"), reverse=True)
        return s_inds[0]
        
    # 从population中选择两个个体
    def Select(self, population, k):
        s_inds = sorted(population, key = itemgetter("fitness"), reverse=True)
        sum_fits = sum(ind['fitness'] for ind in population)
        
        chosen = []
        for i in range(k):
            u = random.random() * sum_fits
            sum_ = 0
            for ind in s_inds:
                sum_ += ind['fitness']
                if sum_ > u:
                    chosen.append(ind)
                    break
        return chosen

    # 交叉繁殖
    def Crossover(self, offspring):
        geninfo1 = offspring[0]['Chromosome']
        geninfo2 = offspring[1]['Chromosome']

        pos1 = random.randrange(1, self.lenchrom)
        pos2 = random.randrange(1, self.lenchrom)

        newoff = []
        for i in range(self.lenchrom):
            if (i >= min(pos1, pos2) and i <= max(pos1, pos2)):
                newoff.append(geninfo2[i])
            else:
                newoff.append(geninfo1[i])
        
        return newoff

    # 变异
    def mutation(self, crossoff):
        pos = random.randrange(1, self.lenchrom)

        if crossoff[pos] == 1:
            crossoff[pos] = 0
        else:
            crossoff[pos] = 1
        return crossoff

    # main frame work of GA
    def GA_main(self):
        print("Start of evolution!")

        population = self.Initialization()
        self.pop = population
        self.bestindividual = self.selectBest(self.pop)
        for g in range(self.maxiter):
            print("-- Generation %i --" % g)
            selectpop = self.Select(self.pop, self.sizepop)
            nextoff = []
            while len(nextoff) != self.sizepop:

                offspring = [random.choice(selectpop) for i in range(2)]

                if random.random() < self.pc:
                    crossoff = self.Crossover(offspring)
                    child = crossoff

                    if random.random() < self.pm:
                        muteoff = self.mutation(crossoff)
                        child = muteoff
                        
                    stepFlag, scores = self.Fobj.step_online(child)
                    if stepFlag == 2:
                        child_fitness = np.mean(scores)/10000
                        nextoff.append({'Chromosome':child, 'fitness':child_fitness})
                    else:
                        continue
            self.pop = nextoff
            # Gather all the fitnesses in one list and print the stats
            fits = [ind['fitness'] for ind in self.pop]

            length = len(self.pop)
            mean = sum(fits) / length
            sum2 = sum(x*x for x in fits)
            std = abs(sum2 / length - mean**2)**0.5
            best_ind = self.selectBest(self.pop)
            if best_ind['fitness'] < self.bestindividual['fitness']:
                self.bestindividual = best_ind
            writer.add_scalar("GA/avg fitness", mean, g)
            writer.add_scalar("GA/min fitness", min(fits), g)
            writer.add_scalar("GA/max fitness", max(fits), g)
            writer.add_scalar("GA/std of pop fitness", std, g)
            writer.add_scalar("GA/best inidival fitness", best_ind['fitness'], g)

        print("-- End of (successful) evolution --")